from docx import Document
import pandas as pd
import os
import sys
from datetime import datetime


def extract_table_content(table):
    """
    提取表格内容为二维列表
    :param table: docx表格对象
    :return: 表格内容列表
    """
    content = []
    for row in table.rows:
        row_content = []
        for cell in row.cells:
            # 提取单元格文本并保留换行符
            cell_text = "\n".join([para.text for para in cell.paragraphs])
            row_content.append(cell_text.strip())
        content.append(row_content)
    return content


def process_documents(doc_zh_path, doc_en_path, output_base_dir):
    """
    处理文档并提取表格内容
    :return: 输出文件路径
    """
    # 检查文件是否存在
    if not os.path.exists(doc_zh_path):
        raise FileNotFoundError(f"中文文档不存在: {doc_zh_path}")
    if not os.path.exists(doc_en_path):
        raise FileNotFoundError(f"英文文档不存在: {doc_en_path}")

    # 提取文档名称（不带扩展名）
    en_doc_name = os.path.splitext(os.path.basename(doc_en_path))[0]
    zh_doc_name = os.path.splitext(os.path.basename(doc_zh_path))[0]

    # 创建输出目录（英文在前，中文在后）
    output_dir = os.path.join(output_base_dir, f"表格对照_{en_doc_name}_{zh_doc_name}")
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 加载文档
    doc_zh = Document(doc_zh_path)
    doc_en = Document(doc_en_path)

    tables_zh = doc_zh.tables
    tables_en = doc_en.tables

    # 确定要处理的表格数量
    table_count = min(len(tables_zh), len(tables_en))
    if table_count == 0:
        print("警告: 未找到任何表格")
        return None

    # 创建输出文件名（英文在前，中文在后）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_file = os.path.join(output_dir, f"{en_doc_name}_{zh_doc_name}_{timestamp}.xlsx")

    # 创建Excel写入器
    with pd.ExcelWriter(output_file, engine='xlsxwriter') as writer:
        # 添加摘要工作表
        summary_data = {
            '属性': ['英文文档', '中文文档', '处理时间', '总表格数'],
            '值': [
                os.path.basename(doc_en_path),
                os.path.basename(doc_zh_path),
                datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                table_count
            ]
        }
        summary_df = pd.DataFrame(summary_data)
        summary_df.to_excel(writer, sheet_name='文档摘要', index=False)

        # 处理每个表格
        for idx in range(table_count):
            try:
                # 提取表格内容
                zh_content = extract_table_content(tables_zh[idx])
                en_content = extract_table_content(tables_en[idx])

                # 计算最大行数和列数
                max_rows = max(len(zh_content), len(en_content))
                max_cols = max(
                    max(len(row) for row in zh_content) if zh_content else 0,
                    max(len(row) for row in en_content) if en_content else 0
                )

                # 创建左右对照的数据结构（英文在前，中文在后）
                table_data = []

                # 添加表头（英文在前）
                table_data.append(["行", "列", "英文内容", "中文内容"])

                # 逐行添加单元格内容
                for i in range(max_rows):
                    # 添加行分隔标记
                    table_data.append([f"--- 第 {i + 1} 行 ---", "", "", ""])

                    # 确定当前行的列数
                    current_cols = max(
                        len(zh_content[i]) if i < len(zh_content) else 0,
                        len(en_content[i]) if i < len(en_content) else 0
                    )

                    # 处理当前行的每个单元格（英文在前）
                    for j in range(max_cols):
                        # 获取英文单元格内容
                        en_cell = ""
                        if i < len(en_content) and j < len(en_content[i]):
                            en_cell = en_content[i][j]

                        # 获取中文单元格内容
                        zh_cell = ""
                        if i < len(zh_content) and j < len(zh_content[i]):
                            zh_cell = zh_content[i][j]

                        # 只添加有内容的行
                        if zh_cell or en_cell:
                            table_data.append([i + 1, j + 1, en_cell, zh_cell])

                # 转换为DataFrame（英文在前）
                combined_df = pd.DataFrame(table_data, columns=['行', '列', '英文', '中文'])

                # 写入工作表
                sheet_name = f"表格_{idx + 1}"
                if len(sheet_name) > 31:  # Excel工作表名称最大长度
                    sheet_name = f"Table_{idx + 1}"

                combined_df.to_excel(writer, sheet_name=sheet_name, index=False, header=False)

                print(f"✓ 已处理表格 {idx + 1}/{table_count}")

            except Exception as e:
                print(f"⚠ 处理表格 {idx + 1} 时出错: {str(e)}")
                continue

    return output_file


if __name__ == "__main__":
    # 示例文件路径 - 在实际使用中替换为您的文件路径
    doc_zh_path = "./test/历史、重组及公司架构/5 中文招股书——历史、重组及公司架构.docx"
    doc_en_path = "./test/历史、重组及公司架构/5 英文招股书——历史、重组及公司架构.docx"
    output_base_dir = "./output"

    # 处理文档
    try:
        output_file = process_documents(doc_zh_path, doc_en_path, output_base_dir)
        if output_file:
            print(f"\n✅ 处理完成! 结果已保存到: {output_file}")
            print(f"共提取 {pd.read_excel(output_file, sheet_name='文档摘要').iloc[3, 1]} 个表格")
        else:
            print("❌ 处理失败，未生成输出文件")
    except Exception as e:
        print(f"❌ 处理过程中出错: {str(e)}")
        sys.exit(1)